"""
检查 NPU 是否可用的测试脚本
"""

import sys


def check_cann_driver():
    """检查昇腾 CANN 驱动是否安装"""
    print("=" * 80)
    print("1. 检查昇腾 CANN 驱动")
    print("=" * 80)
    
    try:
        import subprocess
        
        # 检查 npu-smi 命令（昇腾 NPU 管理工具）
        try:
            result = subprocess.run(['npu-smi', 'info'], 
                                  capture_output=True, 
                                  text=True, 
                                  timeout=5)
            if result.returncode == 0:
                print("✅ CANN 驱动已安装")
                print("\nNPU 设备信息:")
                print(result.stdout[:500])  # 只显示前 500 字符
                return True
            else:
                print("❌ npu-smi 命令失败")
                return False
        except FileNotFoundError:
            print("❌ 未找到 npu-smi 命令")
            print("   说明: CANN 驱动未安装或未配置环境变量")
            print("   安装: 请从华为昇腾官网下载 CANN 工具包")
            return False
        except Exception as e:
            print(f"❌ 检查失败: {e}")
            return False
    except Exception as e:
        print(f"❌ 检查异常: {e}")
        return False


def check_onnxruntime():
    """检查 ONNX Runtime 版本和可用 Providers"""
    print("\n" + "=" * 80)
    print("2. 检查 ONNX Runtime")
    print("=" * 80)
    
    try:
        import onnxruntime as ort
        
        # 检查版本
        version = ort.__version__
        print(f"✅ ONNX Runtime 版本: {version}")
        
        # 检查可用的 Execution Providers
        providers = ort.get_available_providers()
        print(f"\n可用的 Execution Providers ({len(providers)}):")
        
        has_npu = False
        has_gpu = False
        
        for idx, provider in enumerate(providers, 1):
            # 标记当前默认使用的 Provider
            marker = "✅" if idx == 1 else "  "
            print(f"  {marker} {provider}")
            
            # 检查是否有 NPU/GPU
            if 'NPU' in provider.upper() or 'CANN' in provider.upper():
                has_npu = True
            if 'CUDA' in provider.upper() or 'GPU' in provider.upper():
                has_gpu = True
        
        # 总结
        print("\n支持情况:")
        print(f"  NPU 支持: {'✅ 是' if has_npu else '❌ 否'}")
        print(f"  GPU 支持: {'✅ 是' if has_gpu else '❌ 否'}")
        print(f"  CPU 支持: ✅ 是")
        
        if not has_npu:
            print("\n⚠️ 当前 ONNX Runtime 不支持 NPU")
            print("   安装 NPU 版本:")
            print("   1. pip uninstall onnxruntime")
            print("   2. pip install onnxruntime-npu")
        
        return has_npu
        
    except ImportError:
        print("❌ ONNX Runtime 未安装")
        print("   安装: pip install onnxruntime")
        return False
    except Exception as e:
        print(f"❌ 检查失败: {e}")
        return False


def check_rapidocr():
    """检查 RapidOCR 是否可用"""
    print("\n" + "=" * 80)
    print("3. 检查 RapidOCR")
    print("=" * 80)
    
    try:
        from rapidocr_onnxruntime import RapidOCR
        print("✅ RapidOCR 已安装")
        
        # 尝试初始化
        try:
            ocr = RapidOCR()
            print("✅ RapidOCR 初始化成功")
            return True
        except Exception as e:
            print(f"❌ RapidOCR 初始化失败: {e}")
            return False
            
    except ImportError:
        print("❌ RapidOCR 未安装")
        print("   安装: pip install rapidocr-onnxruntime")
        return False


def check_system_info():
    """显示系统信息"""
    print("\n" + "=" * 80)
    print("4. 系统信息")
    print("=" * 80)
    
    print(f"Python 版本: {sys.version}")
    print(f"平台: {sys.platform}")
    
    # 检查环境变量
    import os
    if 'ASCEND_HOME' in os.environ:
        print(f"✅ ASCEND_HOME: {os.environ['ASCEND_HOME']}")
    else:
        print("⚠️ ASCEND_HOME 环境变量未设置")


def test_npu_inference():
    """测试 NPU 推理"""
    print("\n" + "=" * 80)
    print("5. 测试 NPU 推理")
    print("=" * 80)
    
    try:
        from rapidocr_onnxruntime import RapidOCR
        from PIL import Image, ImageDraw, ImageFont
        import time
        
        # 创建测试图片
        print("创建测试图片...")
        img = Image.new('RGB', (400, 100), color='white')
        draw = ImageDraw.Draw(img)
        try:
            font = ImageFont.truetype("simhei.ttf", 30)
        except:
            font = ImageFont.load_default()
        draw.text((20, 35), "测试NPU", fill='black', font=font)
        img.save('test_npu.jpg')
        
        # 测试推理
        print("开始 OCR 识别...")
        ocr = RapidOCR()
        
        start_time = time.time()
        result, elapse = ocr('test_npu.jpg')
        total_time = time.time() - start_time
        
        print(f"✅ 识别完成")
        print(f"  - 耗时: {total_time:.3f} 秒")
        print(f"  - OCR 内部耗时: {elapse:.3f} 秒")
        
        if result:
            texts = [item[1] for item in result]
            print(f"  - 识别结果: {' '.join(texts)}")
        
        # 清理
        import os
        os.remove('test_npu.jpg')
        
        return True
        
    except Exception as e:
        print(f"❌ 推理测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """主检查流程"""
    print("=" * 80)
    print("NPU 环境检查工具")
    print("=" * 80)
    
    results = []
    
    # 1. 检查 CANN 驱动
    has_cann = check_cann_driver()
    results.append(("CANN 驱动", has_cann))
    
    # 2. 检查 ONNX Runtime
    has_npu_runtime = check_onnxruntime()
    results.append(("ONNX Runtime NPU", has_npu_runtime))
    
    # 3. 检查 RapidOCR
    has_rapidocr = check_rapidocr()
    results.append(("RapidOCR", has_rapidocr))
    
    # 4. 系统信息
    check_system_info()
    
    # 5. 测试推理
    if has_rapidocr:
        inference_ok = test_npu_inference()
        results.append(("NPU 推理", inference_ok))
    
    # 总结
    print("\n" + "=" * 80)
    print("检查总结")
    print("=" * 80)
    
    for name, result in results:
        status = "✅ 通过" if result else "❌ 失败"
        print(f"  {name}: {status}")
    
    # 判断 NPU 是否可用
    print("\n" + "=" * 80)
    if has_cann and has_npu_runtime:
        print("✅ NPU 环境已就绪！")
        print("\n启动命令:")
        print("  python file_parse/rapidocr.py --device npu --port 8868")
    elif has_npu_runtime:
        print("⚠️ ONNX Runtime 支持 NPU，但 CANN 驱动可能未安装")
        print("   请确认昇腾驱动是否正确安装")
    elif has_cann:
        print("⚠️ CANN 驱动已安装，但 ONNX Runtime 不支持 NPU")
        print("\n安装 NPU 版 ONNX Runtime:")
        print("  pip uninstall onnxruntime")
        print("  pip install onnxruntime-npu")
    else:
        print("❌ NPU 环境未就绪")
        print("\n配置步骤:")
        print("  1. 安装 CANN 驱动（从华为昇腾官网下载）")
        print("  2. pip install onnxruntime-npu")
        print("  3. pip install rapidocr-onnxruntime")
    
    print("=" * 80)
    
    return has_cann and has_npu_runtime


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
